import numpy as np
import matplotlib.pyplot as plt
import os

# 设置更好的图表风格
plt.style.use('ggplot')

# 加载数据
data_path = './dataset/channel_data/channel_data.npz'
print(f"Loading data: {data_path}")
data = np.load(data_path)

# 获取信道特征数据
channel_features = data['channel_features']  # [n_windows, window_size, n_clusters, 5]
print(f"Channel features shape: {channel_features.shape}")

# 创建保存目录
save_dir = './dataset/channel_data/cluster_metrics'
os.makedirs(save_dir, exist_ok=True)

# 特征名称和单位
feature_names = [
    'Horizontal Arrival Angle (degrees)', 
    'Vertical Arrival Angle (degrees)', 
    'Received Power (dB)', 
    'Horizontal Spread Angle (ASA)', 
    'Vertical Spread Angle (ZSA)'
]

# 为每个窗口创建一个可视化
for window_idx in range(min(10, channel_features.shape[0])):  # 只取前10个窗口
    window_data = channel_features[window_idx]  # [20, 25, 5]
    
    # 获取第一个簇(索引0)在这个窗口的所有时间点上的数据
    cluster_data = window_data[:, 0, :]  # [20, 5]
    
    # 创建图表
    plt.figure(figsize=(15, 12))
    
    # 为每个特征创建一个子图
    for feature_idx in range(5):
        plt.subplot(3, 2, feature_idx+1)
        
        # 获取特征数据
        feature_data = cluster_data[:, feature_idx]
        
        # 绘制折线图
        plt.plot(feature_data, 'o-', linewidth=2, markersize=6)
        plt.title(f'First Cluster: {feature_names[feature_idx]}')
        plt.xlabel('Timestep Index')
        plt.ylabel(feature_names[feature_idx])
        plt.grid(True)
        
        # 添加数据标签
        for i, v in enumerate(feature_data):
            plt.text(i, v, f'{v:.2f}', ha='center', va='bottom', fontsize=8)
    
    # 为整个窗口内的所有特征创建一个归一化对比图
    plt.subplot(3, 2, 6)
    
    # 归一化每个特征以便在同一图表中对比
    normalized_data = np.zeros_like(cluster_data, dtype=float)
    for feature_idx in range(5):
        feature_data = cluster_data[:, feature_idx]
        min_val = feature_data.min()
        max_val = feature_data.max()
        if max_val > min_val:
            normalized_data[:, feature_idx] = (feature_data - min_val) / (max_val - min_val)
        else:
            normalized_data[:, feature_idx] = feature_data - min_val
    
    # 绘制所有特征的归一化曲线
    for feature_idx in range(5):
        plt.plot(normalized_data[:, feature_idx], 'o-', linewidth=2, label=feature_names[feature_idx].split('(')[0])
    
    plt.title('Normalized Metrics for First Cluster')
    plt.xlabel('Timestep Index')
    plt.ylabel('Normalized Value')
    plt.legend()
    plt.grid(True)
    
    # 添加总标题
    plt.suptitle(f'Metrics for First Cluster - Window {window_idx+1}', fontsize=16)
    plt.tight_layout(rect=[0, 0, 1, 0.97])
    
    # 保存图表
    save_path = os.path.join(save_dir, f'first_cluster_window_{window_idx+1}.png')
    plt.savefig(save_path, dpi=150)
    plt.close()
    
    print(f"Saved visualization for window {window_idx+1} to {save_path}")

# 创建一个额外的图表，展示多个窗口中第一个簇的接收功率变化
plt.figure(figsize=(12, 6))

# 选取前5个窗口
for window_idx in range(min(5, channel_features.shape[0])):
    # 获取接收功率 (特征索引2)
    power_data = channel_features[window_idx, :, 0, 2]  # [20]
    plt.plot(power_data, 'o-', linewidth=2, label=f'Window {window_idx+1}')

plt.title('Received Power of First Cluster Across Multiple Windows')
plt.xlabel('Timestep Index')
plt.ylabel('Received Power (dB)')
plt.legend()
plt.grid(True)

# 保存图表
power_save_path = os.path.join(save_dir, 'first_cluster_power_across_windows.png')
plt.savefig(power_save_path, dpi=150)
plt.close()

print(f"Saved multi-window power comparison to {power_save_path}")

# 创建一个统计图表，显示第一个簇的所有特征在全部数据上的分布
plt.figure(figsize=(15, 10))

# 遍历所有特征
for feature_idx in range(5):
    plt.subplot(2, 3, feature_idx+1)
    
    # 获取特定特征在所有窗口、所有时间点上的第一个簇的数据
    feature_values = channel_features[:, :, 0, feature_idx].flatten()
    
    # 绘制直方图
    plt.hist(feature_values, bins=30, alpha=0.7)
    plt.title(f'Distribution of {feature_names[feature_idx].split("(")[0]}')
    plt.xlabel(feature_names[feature_idx])
    plt.ylabel('Frequency')
    plt.grid(True)
    
    # 绘制特征值的统计信息
    plt.text(0.05, 0.95, 
             f'Min: {feature_values.min():.2f}\nMax: {feature_values.max():.2f}\nMean: {feature_values.mean():.2f}\nStd: {feature_values.std():.2f}',
             transform=plt.gca().transAxes,
             verticalalignment='top',
             bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

plt.suptitle('Distribution of Metrics for First Cluster Across All Windows', fontsize=16)
plt.tight_layout(rect=[0, 0, 1, 0.97])

# 保存图表
dist_save_path = os.path.join(save_dir, 'first_cluster_metrics_distribution.png')
plt.savefig(dist_save_path, dpi=150)
plt.close()

print(f"Saved metrics distribution to {dist_save_path}")
print("All visualizations complete.") 